import numpy as np
import pandas as pd
import json
import argparse
import csv
import random
from functools import reduce

max = 210000
num = 10
count = 1000

parser = argparse.ArgumentParser(description='Create test')
parser.add_argument('--mode', default='rw',
                    help='Read/Write Mode.')
args = parser.parse_args()

#args.mode = 'r'

if args.mode == 'rw' or args.mode == 'w':
    # w生成shape写入csv start------------------

    # 数据生成
    # 16的倍数
    def multiples16():
        m16_list = []
        sum = 0
        for i in range(num):
            i = i * sum
            sum = sum + 16
            m16_list.append(sum)
        return m16_list

    # 32的倍数
    def multiples32():
        m32_list = []
        sum = 0
        for i in range(num):
            i = i * sum
            sum = sum + 32
            m32_list.append(sum)
        return m32_list

    # 质数
    def primes():
        primes_list = []
        for i in range(2, num):
            for j in range(2, i):
                if i % j == 0:
                    break
            else:
                primes_list.append(i)
        return primes_list

    # 非16的倍数
    def non16():
        non16_list = []
        for i in range(1, num):
            if i % 16 == 0:
                continue
            else:
                non16_list.append(i)
        return non16_list

    print("===16的倍数===")
    mul16 = multiples16()
    print(mul16)
    print("===32的倍数===")
    mul32 = multiples32()
    print(mul32)
    print("===质数===")
    primes = primes()
    print(primes)
    print("===非16的倍数===")
    non16 = non16()
    print(non16)

    # 一维:取32的倍数
    print("===一维===")
    data1 = []
    for a in range(count):
        data = []
        m32 = random.choice(mul32)
        data.append(m32)
        data1.append(data)
    print(data1[:3])

    print("===二维===")
    # 二维：两个tensor都取32的倍数
    data2 = []
    for i in range(count):
        data = []
        for j in range(2):
            m32 = random.choice(mul32)
            data.append(m32)
        # 列表的乘积
        ln = reduce(lambda x, y: x * y, data)
        if ln > max:  # 如果列表的乘积大于21亿，则退出本次循环，该数据不添加
            continue
        else:
            data2.append(data)
    print(data2[:3])

    print("===三维===")
    # 三维：两个tensor取32的倍数，一个取16的倍数
    data3 = []
    for i in range(count):
        data = []
        for j in range(2):
            m32 = random.choice(mul32)
            data.append(m32)
        m16 = random.choice(mul16)
        data.append(m16)
        ln = reduce(lambda x, y: x * y, data)
        if ln > max:  # 如果列表的乘积大于21亿，则退出本次循环，该数据不添加
            continue
        else:
            data3.append(data)
    print(data3[:3])

    print("===四维===")
    # 四维：一个取16的倍数，一个取32的倍数，一个取质数，一个取非16的倍数
    data4 = []
    for i in range(count):
        data = []
        m16 = random.choice(mul16)
        data.append(m16)
        m32 = random.choice(mul32)
        data.append(m32)
        pr = random.choice(primes)
        data.append(pr)
        n16 = random.choice(non16)
        data.append(n16)
        ln = reduce(lambda x, y: x * y, data)
        if ln > max:  # 如果列表的乘积大于21亿，则退出本次循环，该数据不添加
            continue
        else:
            data4.append(data)
    print(data4[:3])

    print("===五维===")
    # 五维：一个取质数，两个取非16的倍数,一个tensor取16的倍数，一个取32的倍数
    data5 = []
    for i in range(count):
        data = []
        pr = random.choice(primes)
        data.append(pr)
        for i in range(2):
            n16 = random.choice(non16)
            data.append(n16)
        m16 = random.choice(mul16)
        data.append(m16)
        m32 = random.choice(mul32)
        data.append(m32)
        ln = reduce(lambda x, y: x * y, data)
        if ln > max:  # 如果列表的乘积大于21亿，则退出本次循环，该数据不添加
            continue
        else:
            data5.append(data)
    print(data5[:3])

    print("===六维===")
    # 六维：一个tensor取16的倍数，两个取质数，三个取非16的倍数
    data6 = []
    for i in range(count):
        data = []
        m16 = random.choice(mul16)
        data.append(m16)
        for i in range(2):
            pr = random.choice(primes)
            data.append(pr)
        for i in range(3):
            n16 = random.choice(non16)
            data.append(n16)
        ln = reduce(lambda x, y: x * y, data)
        if ln > max:  # 如果列表的乘积大于21亿，则退出本次循环，该数据不添加
            continue
        else:
            data6.append(data)
    print(data6[:3])

    print("===七维===")
    # 七维：三个取质数，一个tensor取16的倍数，三个取非16的倍数
    data7 = []
    for i in range(count):
        data = []
        for i in range(3):
            pr = random.choice(primes)
            data.append(pr)
        m16 = random.choice(mul16)
        data.append(m16)
        for i in range(3):
            n16 = random.choice(non16)
            data.append(n16)
        ln = reduce(lambda x, y: x * y, data)
        if ln > max:  # 如果列表的乘积大于21亿，则退出本次循环，该数据不添加
            continue
        else:
            data7.append(data)
    print(data7[:3])

    print("===八维===")
    # 八维：三个取质数，五个取非16的倍数
    data8 = []
    for i in range(count):
        data = []
        for i in range(3):
            pr = random.choice(primes)
            data.append(pr)
        for i in range(5):
            n16 = random.choice(non16)
            data.append(n16)
        ln = reduce(lambda x, y: x * y, data)
        if ln > max:  # 如果列表的乘积大于21亿，则退出本次循环，该数据不添加
            continue
        else:
            data8.append(data)
    print(data8[:3])


    # 写入csv(根据算子需要进行修改，也可在csv生成后直接修改csv)
    rows = [
        ['op', 'inputs', '', '', '', '', 'outputs', '', '', '', '', 'attrs', '', ''],
        ['Abs', 'shape', 'format', 'name', 'type', 'value_range', 'shape', 'format', 'name', 'type', 'value_range', 'name',
         'type', 'value'],
        ['', data1[:3], 'ND', 'input_x;window', 'float32', [-0.1, 0], data1[:3], 'ND', 'output_x', 'float32', [-0.1, 0], '',
         '', ''],
        ['', data2[:3], 'ND', 'input_x', 'float16', [-1, 0], data2[:3], 'ND', 'output_x', 'float16', [-1, 0], '',
         '', ''],
        ['', data3[:3], 'ND', 'input_x', 'int32', [-100, 0], data3[:3], 'ND', 'output_x', 'int32', [-100, 0], '',
         '', ''],
        ['', data4[:3], 'ND', 'input_x', 'float32', [0, 0.1], data4[:3], 'ND', 'output_x', 'float32', [0, 0.1], '',
         '', ''],
        ['', data5[:3], 'ND', 'input_x', 'float32', [0, 1], data5[:3], 'ND', 'output_x', 'float32', [0, 1], '',
         '', ''],
        ['', data6[:3], 'ND', 'input_x', 'int64', [0, 100], data6[:3], 'ND', 'output_x', 'int64', [0, 100], '',
         '', ''],
        ['', data7[:3], 'ND', 'input_x', 'int32', [-10, 10], data7[:3], 'ND', 'output_x', 'int32', [-10, 10], '',
         '', ''],
        ['', data8[:3], 'ND', 'input_x', 'float32', [-0.1, 0.1], data8[:3], 'ND', 'output_x', 'float32', [0, 10], '',
         '', ''],
    ]

    with open('test_csv.csv', 'w+', newline='')as f:
        f_csv = csv.writer(f)
        f_csv.writerows(rows)
# w生成shape写入csv end------------------


# r读取csv生成st start------------------
# str:"[123]"转换为list:[123]
def str2list(shape_string):
    shape_split = shape_string[1:-1].split(",")
    shape_split = list(map(int, shape_split))
    return shape_split
# str:"[[123],[456]]"转换为list: [[123],[456]]
def str2list2(shape_string):
    out = []
    shape_split = shape_string[2:-2].split("], [")
    for i in range(len(shape_split)):
        a = shape_split[i].split(',')
        b = list(map(int, a))
        out.append(b)
    return out
# str:"[123]"转换为list:[123]
def str2list_float(shape_string):
    shape_split = shape_string[1:-1].split(",")
    shape_split = list(map(float, shape_split))
    return shape_split
# str:"[[123],[456]]"转换为list: [[123],[456]]
def str2list2_float(shape_string):
    out = []
    shape_split = shape_string[2:-2].split("], [")
    for i in range(len(shape_split)):
        a = shape_split[i].split(',')
        b = list(map(float, a))
        out.append(b)
    return out
if args.mode == 'rw' or args.mode == 'r':
    csv_reader = pd.read_csv("test_csv.csv", encoding='utf-8', header=None)

    json_list = []
    json_dict = dict()
    for test_idex in range(2, csv_reader.shape[0]):
        json_dict = dict()
        json_dict["case_name"] = "Test_" + csv_reader[0][1] + "_case" + str(test_idex - 1)
        json_dict["op"] = csv_reader[0][1]
        json_dict["calc_expect_func_file"] = ""
        inputs_list = list()
        json_dict["input_desc"] = inputs_list
        outputs_list = list()
        json_dict["output_desc"] = outputs_list
        attrs_list = list()
        json_dict["attr"] = attrs_list

        type_list = list(csv_reader.iloc[0])
        inputs_index = type_list.index("inputs")
        outputs_index = type_list.index("outputs")
        attrs_index = type_list.index("attrs")
        type_list_length = len(type_list)

        # 读csv中算子输入
        inputs_num = len(csv_reader[inputs_index][test_idex].split(";"))
        for i in range(inputs_num):
            inputs_list.append(dict())
        for i in range(inputs_index, outputs_index):
            input_info_list = csv_reader[i][test_idex].split(";")
            for j in range(inputs_num):
                if "shape" in csv_reader[i][1]:
                    if '[[' in input_info_list[j]:
                        inputs_list[j][csv_reader[i][1]] = str2list2(input_info_list[j])
                    else:
                        inputs_list[j][csv_reader[i][1]] = str2list(input_info_list[j])
                if "format" in csv_reader[i][1]:
                    inputs_list[j][csv_reader[i][1]] = input_info_list[j]
                if "name" in csv_reader[i][1]:
                    inputs_list[j][csv_reader[i][1]] = input_info_list[j]
                if "type" in csv_reader[i][1]:
                    inputs_list[j][csv_reader[i][1]] = input_info_list[j]
                if "value_range" in csv_reader[i][1]:
                    if '[[' in input_info_list[j]:
                        inputs_list[j][csv_reader[i][1]] = str2list2_float(input_info_list[j])
                    else:
                        inputs_list[j][csv_reader[i][1]] = str2list_float(input_info_list[j])
                if "value" in csv_reader[i][1]:
                    try:
                        if '[[' in input_info_list[j]:
                            inputs_list[j][csv_reader[i][1]] = str2list2_float(input_info_list[j])
                        else:
                            inputs_list[j][csv_reader[i][1]] = str2list_float(input_info_list[j])
                    except ValueError:
                        inputs_list[j][csv_reader[i][1]] = []

        # 读csv中算子输出
        outputs_num = len(csv_reader[outputs_index][test_idex].split(";"))
        for i in range(outputs_num):
            outputs_list.append(dict())
        for i in range(outputs_index, attrs_index):
            output_info_list = csv_reader[i][test_idex].split(";")
            for j in range(outputs_num):
                if "shape" in csv_reader[i][1]:
                    if '[[' in output_info_list[j]:
                        outputs_list[j][csv_reader[i][1]] = str2list2(output_info_list[j])
                    else:
                        outputs_list[j][csv_reader[i][1]] = str2list(output_info_list[j])
                if "format" in csv_reader[i][1]:
                    outputs_list[j][csv_reader[i][1]] = output_info_list[j]
                if "name" in csv_reader[i][1]:
                    outputs_list[j][csv_reader[i][1]] = output_info_list[j]
                if "type" in csv_reader[i][1]:
                    outputs_list[j][csv_reader[i][1]] = output_info_list[j]
                if "value_range" in csv_reader[i][1]:
                    outputs_list[j][csv_reader[i][1]] = output_info_list[j]
                    if '[[' in output_info_list[j]:
                        outputs_list[j][csv_reader[i][1]] = str2list2_float(output_info_list[j])
                    else:
                        outputs_list[j][csv_reader[i][1]] = str2list_float(output_info_list[j])
        # 读csv中算子属性
        if not isinstance(csv_reader[attrs_index][test_idex],float):
            attrs_num = len(csv_reader[attrs_index][test_idex].split(";"))
            for i in range(attrs_num):
                attrs_list.append(dict())
            for i in range(attrs_index, type_list_length):
                attr_info_list = csv_reader[i][test_idex].split(";")
                for j in range(attrs_num):
                    if "value" in csv_reader[i][1]:
                        if attrs_list[j]["type"] == "string":
                            attrs_list[j][csv_reader[i][1]] = attr_info_list[j]
                        elif attrs_list[j]["type"] == "bool":
                            attrs_list[j][csv_reader[i][1]] = True if attr_info_list[j] == "true" else False
                        elif attrs_list[j]["type"] == "list_int":
                            attrs_list[j][csv_reader[i][1]] = str2list(attr_info_list[j])
                        elif attrs_list[j]["type"] == "float" or attrs_list[j]["type"] == "float16":
                            attrs_list[j][csv_reader[i][1]] = float(attr_info_list[j])
                        else:
                            attrs_list[j][csv_reader[i][1]] = int(attr_info_list[j])
                    else:
                        attrs_list[j][csv_reader[i][1]] = attr_info_list[j]

        json_list.append(json_dict)

    # 写入csv内容,生成st
    st_file = "test_" + json_dict["op"] + "_case.json"
    with open(st_file, "w") as f:
        f.write(json.dumps(obj=json_list, ensure_ascii=True))
        f.close()
# r读取csv生成st end------------------
